// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <optional>
#include <pollux/testing/gtest_utils.h>
#include <pollux/common/testutil/optional_empty.h>
#include <pollux/testing/functions/prestosql/function_base_test.h>

using namespace kumo::pollux::test;

namespace kumo::pollux::functions {
namespace {
class ArrayNGramsTest : public test::FunctionBaseTest {
 protected:
  template <typename T>
  void testNgram(
      const std::vector<std::optional<T>>& inputArray,
      int32_t n,
      const std::vector<std::optional<std::vector<std::optional<T>>>>&
          expectedOutput) {
    std::vector<std::optional<std::vector<std::optional<T>>>> inputVec(
        {inputArray});
    auto input = make_nullable_array_vector<T>(inputVec);
    auto result = evaluate(
        fmt::format("ngrams(c0, {}::INTEGER)", n), make_row_vector({input}));

    auto expected = make_nullable_nested_array_vector<T>({expectedOutput});
    assertEqualVectors(expected, result);
  }
};

TEST_F(ArrayNGramsTest, integers) {
  testNgram<int64_t>({1, 2, 3, 4}, 1, {{{1}}, {{2}}, {{3}}, {{4}}});
  testNgram<int64_t>({1, 2, 3, 4}, 2, {{{1, 2}}, {{2, 3}}, {{3, 4}}});
  testNgram<int64_t>({1, 2, 3, 4}, 3, {{{1, 2, 3}}, {{2, 3, 4}}});
  testNgram<int64_t>({1, 2, 3, 4}, 4, {{{1, 2, 3, 4}}});
  testNgram<int64_t>({1, 2, 3, 4}, 5, {{{1, 2, 3, 4}}});
  testNgram<int64_t>(
      {1, 2, 3, 4}, std::numeric_limits<int32_t>::max(), {{{1, 2, 3, 4}}});
  testNgram<int64_t>({}, 1, {common::testutil::optionalEmpty});
  testNgram<int64_t>({}, 10, {common::testutil::optionalEmpty});
}

TEST_F(ArrayNGramsTest, invalidN) {
  auto input = make_array_vector<int64_t>({{1, 2, 3, 4}});
  POLLUX_ASSERT_THROW(
      evaluate("ngrams(c0, 0::INTEGER)", make_row_vector({input})),
      "(0 vs. 0) N must be greater than zero");
  POLLUX_ASSERT_THROW(
      evaluate("ngrams(c0, -5::INTEGER)", make_row_vector({input})),
      "Scalar function signature is not supported");
  input = make_array_vector<int64_t>({{}});
  POLLUX_ASSERT_THROW(
      evaluate("ngrams(c0, 0::INTEGER)", make_row_vector({input})),
      "(0 vs. 0) N must be greater than zero");
}

TEST_F(ArrayNGramsTest, strings) {
  testNgram<std::string>(
      {"foo", "bar", "baz", "this is a very long sentence"},
      1,
      {{{"foo"}}, {{"bar"}}, {{"baz"}}, {{"this is a very long sentence"}}});
  testNgram<std::string>(
      {"foo", "bar", "baz", "this is a very long sentence"},
      2,
      {{{"foo", "bar"}},
       {{"bar", "baz"}},
       {{"baz", "this is a very long sentence"}}});
  testNgram<std::string>(
      {"foo", "bar", "baz", "this is a very long sentence"},
      3,
      {{{"foo", "bar", "baz"}},
       {{"bar", "baz", "this is a very long sentence"}}});
  testNgram<std::string>(
      {"foo", "bar", "baz", "this is a very long sentence"},
      4,
      {{{"foo", "bar", "baz", "this is a very long sentence"}}});
  testNgram<std::string>(
      {"foo", "bar", "baz", "this is a very long sentence"},
      5,
      {{{"foo", "bar", "baz", "this is a very long sentence"}}});
}

TEST_F(ArrayNGramsTest, nulls) {
  testNgram<std::string>(
      {"foo", std::nullopt, "bar"},
      2,
      {{{"foo", std::nullopt}}, {{std::nullopt, "bar"}}});
  testNgram<std::string>(
      {std::nullopt, std::nullopt, std::nullopt},
      2,
      {{{std::nullopt, std::nullopt}}, {{std::nullopt, std::nullopt}}});
}
} // namespace

} // namespace kumo::pollux::functions
